[DE-8270] Model weights upload & download (SDK side) by luke-e-schaefer · Pull Request #469 · scaleapi/nucleus-python-client

luke-e-schaefer · 2026-07-27T20:36:33Z

Summary

SDK side of model weights upload / download, mirroring the REST surface merged in scaleapi#149063.

The two primary methods are the ones the server PR's API-docs pages (ApiDocsPage/models-python/{upload,download}-model-weights.md) already document, so the published docs and the SDK agree:

import nucleus

client = nucleus.NucleusClient(YOUR_SCALE_API_KEY)
model = client.get_model(reference_id="My-CNN")

client.upload_model_weights(model, "/path/to/weights.bin")
client.download_model_weights(model, "/path/to/save/weights.bin")

Added

NucleusClient.upload_model_weights(model, path, *, content_type=None, original_filename=None, checksum_sha256=None, on_progress=None) — presign → PUT direct to storage → finalize. Returns ModelWeights.
NucleusClient.download_model_weights(model, path, *, on_progress=None) — resolves the signed URL and streams to disk (creating parent dirs). Returns the path written.
get_model_weights(model) / delete_model_weights(model) for the remaining two routes.
Model.upload_weights() / .download_weights() / .weights() / .delete_weights() — thin delegation to the client, matching how Benchmark wraps its client methods.
ModelWeights metadata type: present, status, size_bytes, original_filename, content_type, download_url.

All model arguments accept either a Model or a bare model id (prj_*).

Notes for review

Bytes never transit the Nucleus API. Transfers go straight to storage over presigned URLs, so multi-GB artifacts aren't subject to API request-size limits.
Part PUTs are sent with no headers. They're signed without the Content-Type condition, so forwarding requiredHeaders (which the single PUT does need) makes S3 reject the signature. There's a test pinning this in both directions — it's the easiest thing to get wrong here, and the frontend hook in 149063 has the same split.
Download uses ?json=1 to fetch the signed URL rather than following the 302, so the API's auth headers are never sent to storage.
Multipart above 5 GB, 4 parts in flight (a single S3 PUT is connection-throughput-bound). Missing part ETags fail on the first part rather than after transferring everything and dying at finalize.
The 10 GB server cap is checked client-side before presign, so an oversized file fails without a network round-trip.
The weights routes serialize camelCase in both directions, unlike most endpoints this SDK talks to, so the new payload keys are grouped and labelled as such in constants.py.

Tests / Version

tests/test_model_weights.py — 24 mock-based unit tests (no live API, no real S3): DTO parsing, payload builders, single vs. multipart transfer, header split, ETag/failure handling, progress callbacks, download streaming, all four client methods, and the Model wrappers.
Verified locally the way CI does: pylint nucleus 10.00/10, mypy --ignore-missing-imports nucleus clean, ruff clean, black + isort clean, 61 mock-based tests passing across the eval/benchmark/leaderboard/weights suites.
pyproject.toml → 0.19.1 + CHANGELOG entry (patch bump: additive new methods, per CLAUDE.md).

resolves https://linear.app/scale-epd/issue/DE-8270

🤖 Generated with Claude Code

Greptile Summary

This PR adds model weights upload and download to the Nucleus Python SDK, routing transfers directly to storage via presigned URLs so that multi-GB artifacts never pass through the Nucleus API. The implementation handles both single-PUT (< 5 GB) and S3 multipart (≥ 5 GB) paths with exponential-backoff retries, a threading-lock-protected progress counter for concurrent part uploads, atomic file replacement on download, and a tqdm progress bar that callers can suppress.

NucleusClient gains upload_model_weights, download_model_weights, get_model_weights, and delete_model_weights; Model adds thin delegation wrappers (upload_weights, download_weights, weights, delete_weights) matching the existing Benchmark pattern.
ModelWeights is the new metadata dataclass; nucleus/constants.py adds 14 camelCase wire-format keys (explicitly labelled as such since these routes do not snake-case their DTOs).
24 mock-based unit tests cover the full transfer surface including the per-part header split, ETag failure propagation, concurrency non-overlap, atomic download replacement, and all retry scenarios.

Confidence Score: 5/5

Safe to merge. The transfer logic, retry handling, thread safety, and atomic file replacement are all correctly implemented and well-tested.
The change is additive — four new client methods and one new module, with no modifications to existing code paths. The two minor findings (tqdm bar can show negative increments on a retried single-PUT, and multipart workers run to completion after a part error) are cosmetic/efficiency issues that don't affect data integrity or correctness. The core concerns from previous review rounds (race condition on the transferred counter, single-PUT progress reporting) have been addressed.
nucleus/model_weights.py — specifically the _ProgressReader reset path and ThreadPoolExecutor shutdown behaviour on multipart failure, both noted in review comments.

Important Files Changed

Filename	Overview
nucleus/model_weights.py	New core module implementing presign → PUT/multipart → finalize upload and streaming download with retry, progress reporting, and atomic file replacement. Logic is sound; two minor issues: progress bar can show negative deltas on single-PUT retry (bar.n not reset alongside _ProgressReader._transferred), and all multipart tasks run to completion even after a part fails due to ThreadPoolExecutor pool.map semantics.
nucleus/init.py	Adds upload_model_weights, download_model_weights, get_model_weights, and delete_model_weights to NucleusClient. Imports private helpers (_stream_weights_to_file etc.) from model_weights.py so they can be used inside the class body and patched in tests via nucleus.*. Implementation is clean and consistent with existing client patterns.
nucleus/model.py	Adds four thin delegation methods (upload_weights, download_weights, weights, delete_weights) to Model, exactly mirroring how Benchmark wraps client methods. No logic here.
nucleus/constants.py	Adds 14 camelCase constants for the model-weights wire format. The comment explaining why these are camelCase (unlike the rest of the SDK) is helpful. Constants are accurate representations of their values.
tests/test_model_weights.py	24 well-structured unit tests covering DTO parsing, payload builders, single/multipart transfers, header split, ETag failure, progress callbacks (including the concurrency non-overlap test), retry/backoff, download streaming, atomic file replacement, and all four client methods plus Model wrappers. No live API or S3 calls.

Sequence Diagram

sequenceDiagram
    participant User
    participant NucleusClient
    participant NucleusAPI
    participant Storage

    Note over User,Storage: Upload flow
    User->>NucleusClient: upload_model_weights(model, path)
    NucleusClient->>NucleusClient: os.path.getsize(path) → check ≤ 10 GB
    NucleusClient->>NucleusAPI: "POST model/{id}/weights/presign"
    NucleusAPI-->>NucleusClient: "{uploadId, uploadUrl|parts, requiredHeaders}"

    alt "Single PUT (< 5 GB)"
        NucleusClient->>Storage: PUT uploadUrl (with requiredHeaders, _ProgressReader)
        Storage-->>NucleusClient: ETag
    else Multipart (≥ 5 GB)
        par 4 concurrent workers
            NucleusClient->>Storage: PUT parts[n].url (no headers)
            Storage-->>NucleusClient: ETag[n]
        end
    end

    NucleusClient->>NucleusAPI: "POST model/{id}/weights/finalize {uploadId, parts}"
    NucleusAPI-->>NucleusClient: ModelWeights JSON
    NucleusClient-->>User: ModelWeights

    Note over User,Storage: Download flow
    User->>NucleusClient: download_model_weights(model, path)
    NucleusClient->>NucleusAPI: "GET model/{id}/weights/download?json=1"
    NucleusAPI-->>NucleusClient: "{url: signed-GET-URL}"
    NucleusClient->>Storage: GET signed-URL (streaming → temp file)
    Storage-->>NucleusClient: body chunks
    NucleusClient->>NucleusClient: os.replace(partial_path, path) [atomic]
    NucleusClient-->>User: path written

_{Reviews (11): Last reviewed commit: "refactor(weights): drop the unused `tota..." | Re-trigger Greptile}

Mirrors the REST surface shipped in scaleapi#149063 so users can attach a weights artifact to a model and fetch it back from Python. - `NucleusClient.upload_model_weights` / `download_model_weights` — the two methods the server PR's API-docs pages already document — plus `get_model_weights` and `delete_model_weights` for the remaining routes. - `Model.upload_weights()` / `download_weights()` / `weights()` / `delete_weights()` delegate to the client, matching how `Benchmark` does it. - New `ModelWeights` metadata type parsed from the weights DTO. The weights routes serialize camelCase both ways, unlike most of this SDK's endpoints, so the new payload keys are grouped and labelled in `constants.py`. - Transfers go straight to storage via presigned URLs and never through the API, so artifacts aren't subject to API request-size limits. Over 5 GB the server hands back multipart parts, which upload 4 at a time; `on_progress` reports `(bytes_transferred, total_bytes)`. - Size is checked against the server's 10 GB cap before presign, so an oversized file fails without a network round-trip. Two things worth knowing for review: part PUTs must be sent with *no* headers (they're signed without the Content-Type condition, so forwarding `requiredHeaders` makes S3 reject the signature), and download resolves the signed URL via `?json=1` rather than following the 302, so the API's auth headers are never sent to storage. 24 mock-based unit tests in `tests/test_model_weights.py`; version bumped to 0.19.1 (additive, per CLAUDE.md). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The docstrings are what users read, so they shouldn't describe how the artifact gets stored or moved. Dropped the presign/multipart/direct-to-storage narration from every public docstring, the `ModelWeights` attribute docs, and the CHANGELOG, leaving what a caller actually needs: what the method does, who can call it, the size limit, and the arguments. Also made the transfer helpers private (`_presign_payload`, `_transfer_weights_to_storage`, `_stream_weights_to_file`, `_finalize_payload`) so the mechanics don't show up in the generated API docs at all, rather than only being reworded. Kept the two in-body comments that explain why part uploads send no headers and why the download URL is fetched as JSON — those aren't user-visible and each one guards a real footgun. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…progress Addresses two review comments: - transferred += len(chunk) ran unsynchronized across the part-upload pool, so concurrent workers could drop updates. The counter and the value handed to on_progress are now taken under a lock. - A single PUT reported nothing until it finished, then jumped to 100%. When a callback is supplied the body is wrapped so progress comes from the read side; the wrapper delegates everything but read(), so requests still sizes the body from fileno()/tell() and sends Content-Length as before.

Resolves two version-bump conflicts from #470 (v0.20.0): - pyproject.toml: 0.19.2/0.20.0 -> 0.20.1 - CHANGELOG.md: keep both sections, retitle the weights entry to 0.20.1 nucleus/__init__.py auto-merged cleanly (master touched create_benchmark, this branch adds the model-weights methods). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Three fixes from review: - Download streamed straight into the destination, so an interrupted transfer left a truncated artifact that looked complete. Stream to a sibling temp file and os.replace() on success; a failed re-download now also leaves any existing artifact intact. - The multipart progress counter was locked but on_progress was called after releasing, so two threads could compute 100 and 200 and then call in either order. Invoke the callback under the same lock. - Each in-flight part is read fully into memory, so peak usage was 4 * partSizeBytes with a server-chosen part size. Bound concurrency by a 512 MB budget via _part_upload_workers(), never below 1. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

edwinpav · 2026-08-14T16:54:15Z

👀

edwinpav

overall looks good. mostly nit comments - not too familiar with file reading and writing but the main comments are about error handling around when a non-expected amount of bytes are read/written. also for progress bar, didn't follow it super well but i've used tqdm package in python before which works well (i think it's already used within repo as well)

edwinpav · 2026-08-14T18:06:44Z

+        parts = _transfer_weights_to_storage(
+            path, presign, total_bytes, on_progress
+        )


should this be wrapped in a try/catch?

id probably want to let this bubble up then catch and rethrow imo

that works - so it's caught at some point?

no like it should error out...if this is failing we should know quick iykwim

unless you're concerned about security issues in the error stack

- upload_model_weights: expand `~` in path, hoist the resolved filename to a local, and validate the presign `uploadId` before starting the (up to 10 GB) transfer so a malformed presign fails fast instead of KeyError-ing at finalize. - download_model_weights: expand `~` in path; raise NotFoundError (the codebase idiom) instead of ValueError when a model has no weights artifact. - _stream_weights_to_file: verify received bytes match Content-Length before the atomic replace, so a short/interrupted stream can't promote a truncated file. - Model.upload_weights/download_weights: mirror the client's keyword-only params explicitly instead of **kwargs (better IDE/type-checker support); drop the now redundant quotes on the ModelWeights return annotations. - Tests: assert per-part bytes/offsets; add short-final-part, multipart-through- upload_model_weights (ordered finalize payload), concurrent part-failure, and raising-progress-callback coverage. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…rop unused _client Addresses the remaining review feedback: - Progress is now a tqdm bar shown by default (notebook-aware via the client's tqdm_bar), replacing the raw on_progress callback on the public API. Pass progress=False to silence it. upload_model_weights/download_model_weights and the Model wrappers take `progress: bool = True`. - S3 PUTs and the download GET now retry transient failures (network errors, 429, 5xx) up to 5 attempts with exponential backoff before giving up; 4xx is treated as terminal. Rewindable bodies (single PUT, download temp file) are re-seeked/truncated between attempts so a retry never sends a truncated body. The short-read download check now triggers a retry rather than a hard failure. - Removed the unused ModelWeights._client field and the client arg to from_json. - Tests: retry paths (transient-then-success, network error, exhaustion, no retry on 4xx, single-PUT re-seek, download retry) and tqdm on/off; autouse fixture no-ops backoff sleeps. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

edwinpav · 2026-08-17T21:52:32Z

            checksum_sha256: Optional SHA-256 of the artifact.
-            on_progress: Called with ``(bytes_uploaded, total_bytes)`` as the
-              upload proceeds.
+            progress: Whether to show a ``tqdm`` progress bar for the upload.


nit: any reason why we don't always show this progress bar? just curious

might be loud for really big ones and users could just want a clean console...default is on, but maybe just in case you're wanting to pipe the stdout somewhere else

The tqdm adapter must match the ProgressCallback (transferred, total) signature, but it ignores total (the bar already knows it). Rename to _total and document why, dropping the pylint disable. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

on_progress is internal-only (the public API takes progress: bool), and the sole consumer is the tqdm adapter, which ignores total since the bar already knows it. So ProgressCallback is now Callable[[int], None] — cumulative bytes only — instead of carrying a total that nothing reads. The download completeness check still uses Content-Length directly; it never went through the callback. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

greptile-apps Bot reviewed Jul 27, 2026

View reviewed changes

Comment thread nucleus/model_weights.py

Comment thread nucleus/model_weights.py Outdated

Base automatically changed from update-nuc-sdk-for-new-eval-stuff-pt1 to master August 11, 2026 14:24

luke-e-schaefer and others added 2 commits August 11, 2026 14:31

luke-e-schaefer force-pushed the lukeschaefer/de-8270-upload-download-model-weights branch from df6ede7 to 83c25f7 Compare August 11, 2026 14:35

luke-e-schaefer and others added 2 commits August 11, 2026 15:24

edwinpav self-requested a review August 13, 2026 16:14

edwinpav reviewed Aug 14, 2026

View reviewed changes

luke-e-schaefer and others added 2 commits August 17, 2026 12:17

luke-e-schaefer requested a review from edwinpav August 17, 2026 18:28

edwinpav approved these changes Aug 17, 2026

View reviewed changes

luke-e-schaefer and others added 2 commits August 17, 2026 17:36

luke-e-schaefer merged commit 62e8f8d into master Aug 17, 2026
9 checks passed

luke-e-schaefer deleted the lukeschaefer/de-8270-upload-download-model-weights branch August 17, 2026 23:20

Uh oh!

Conversation

luke-e-schaefer commented Jul 27, 2026 • edited by greptile-apps Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Added

Notes for review

Tests / Version

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

edwinpav commented Aug 14, 2026

Uh oh!

edwinpav left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

edwinpav Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

luke-e-schaefer Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

edwinpav Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

luke-e-schaefer Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

luke-e-schaefer Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

edwinpav Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

luke-e-schaefer Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

luke-e-schaefer commented Jul 27, 2026 •

edited by greptile-apps Bot

Loading